#!/usr/bin/env python3
"""检查 ChromaDB 数据格式"""

import asyncio
import sys
import os
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

from src.config import load_config
from src.database.chromadb_client import ChromaDBClient

async def check_chromadb_data_format():
    """检查 ChromaDB 数据格式"""
    print("=== 检查 ChromaDB 数据格式 ===")
    
    try:
        # 加载配置
        config = load_config()
        
        # 初始化 ChromaDBClient
        chromadb_client = ChromaDBClient(config)
        await chromadb_client.connect()
        
        print("✅ ChromaDBClient 连接成功")
        
        # 获取 conversations 集合
        conversations_collection = chromadb_client.conversations_collection
        
        # 获取所有数据
        print(f"\n1. 获取 conversations 集合中的所有数据...")
        
        # 使用 get 方法获取所有数据
        all_data = conversations_collection.get()
        
        print(f"   总文档数: {len(all_data['ids'])}")
        
        # 分析数据结构
        print(f"\n2. 分析数据结构...")
        
        for i in range(min(len(all_data['ids']), 10)):  # 只显示前10个
            print(f"\n   文档 [{i+1}]:")
            print(f"     ID: {all_data['ids'][i]}")
            print(f"     内容: {all_data['documents'][i][:100]}...")
            
            # 检查元数据
            metadata = all_data['metadatas'][i] if all_data['metadatas'] else {}
            print(f"     元数据:")
            for key, value in metadata.items():
                print(f"       {key}: {value}")
        
        # 检查用户ID分布
        print(f"\n3. 检查用户ID分布...")
        user_ids = {}
        conversation_ids = set()
        
        for i in range(len(all_data['ids'])):
            metadata = all_data['metadatas'][i] if all_data['metadatas'] else {}
            user_id = metadata.get('user_id', 'unknown')
            conv_id = metadata.get('conversation_id', 'unknown')
            
            if user_id not in user_ids:
                user_ids[user_id] = 0
            user_ids[user_id] += 1
            
            conversation_ids.add(conv_id)
        
        print(f"   用户ID分布:")
        for user_id, count in user_ids.items():
            print(f"     {user_id}: {count} 条消息")
        
        print(f"\n   对话ID列表:")
        for conv_id in sorted(conversation_ids):
            print(f"     {conv_id}")
        
        # 搜索测试
        print(f"\n4. 搜索测试...")
        
        # 不带用户过滤的搜索
        print(f"   不带用户过滤搜索 '王者荣耀':")
        results_no_filter = conversations_collection.query(
            query_texts=["王者荣耀"],
            n_results=10
        )
        
        print(f"     结果数量: {len(results_no_filter['ids'][0])}")
        for i, result_id in enumerate(results_no_filter['ids'][0]):
            metadata = results_no_filter['metadatas'][0][i]
            distance = results_no_filter['distances'][0][i]
            print(f"     [{i+1}] ID: {result_id}")
            print(f"         对话ID: {metadata.get('conversation_id')}")
            print(f"         用户ID: {metadata.get('user_id')}")
            print(f"         距离: {distance}")
        
        # 带用户过滤的搜索
        print(f"\n   带用户过滤搜索 '王者荣耀' (user_id=default_user):")
        results_with_filter = conversations_collection.query(
            query_texts=["王者荣耀"],
            n_results=10,
            where={"user_id": "default_user"}
        )
        
        print(f"     结果数量: {len(results_with_filter['ids'][0])}")
        for i, result_id in enumerate(results_with_filter['ids'][0]):
            metadata = results_with_filter['metadatas'][0][i]
            distance = results_with_filter['distances'][0][i]
            print(f"     [{i+1}] ID: {result_id}")
            print(f"         对话ID: {metadata.get('conversation_id')}")
            print(f"         用户ID: {metadata.get('user_id')}")
            print(f"         距离: {distance}")
        
        print(f"\n✅ 检查完成")
        
    except Exception as e:
        print(f"❌ 检查失败: {e}")
        import traceback
        traceback.print_exc()

if __name__ == "__main__":
    asyncio.run(check_chromadb_data_format())